package io.zouyalong.algo.leetcode;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;


/**
 * https://leetcode-cn.com/problems/n-queens/
 */
public class NQueen {
    public static void main(String[] args) {
        int n = 5;
        var inst = new NQueen();
        System.out.printf("Results: %s", inst.solution1(n));
    }

    class Board {
        boolean occupied[][];
        int size;
        Set<Integer> rows;
        Set<Integer> cols;
        Set<Integer> leading; // 主对角线
        Set<Integer> minors;    // 次对角线

        public Board(int n) {
            size = n;
            occupied = new boolean[n][n];
            rows = new HashSet<>();
            cols = new HashSet<>();
            leading = new HashSet<>();
            minors = new HashSet<>();
        }

        public boolean isValid(int row, int col) {
            return !rows.contains(row) && !cols.contains(col)
                    && !leading.contains(row-col)
                    && !minors.contains(row+col);
        }

        public void setOccupied(int row, int col) {
            if (!isValid(row, col)) {
                return;
            }
            occupied[row][col] = true;
            rows.add(row);
            cols.add(col);
            leading.add(row-col);
            minors.add(row+col);
        }

        public void clearOccupied(int row, int col) {
            if (occupied[row][col]) {
                occupied[row][col] = false;
                rows.remove(row);
                cols.remove(col);
                leading.remove(row-col);
                minors.remove(row+col);
            }
        }

        public List<String> convertToString() {
            ArrayList<String> results = new ArrayList<>(this.size);
            for (int i = 0; i < this.size; i++) {
                var chars = new char[this.size];
                for (int j = 0; j < this.size; j++) {
                    if (occupied[i][j]) {
                        chars[j] = 'Q';
                    } else {
                        chars[j] = '.';
                    }
                }
                results.add(String.valueOf(chars));
            }
            return results;
        }
    }

    public List<List<String>> solution1(int n) {
        if (n == 1) {
            return List.of(List.of("Q"));
        }
        var board = new Board(n);
        List<List<String>> results = new ArrayList<>();
        backtrace(results, board, 0);
        return results;
    }

    public void backtrace(List<List<String>> results, Board board, int row) {
        if (row == board.size) {
            results.add(board.convertToString());
            return;
        }
        for (int col = 0; col < board.size; col++) {
            if (!board.isValid(row, col)) continue;
            board.setOccupied(row, col);
            backtrace(results, board, row+1);
            board.clearOccupied(row, col);
        }
    }
}
